package com.le.tester.imooc.oldfairy.dbgen;

import org.apache.commons.lang3.tuple.Pair;
import org.testng.annotations.Test;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;



/**
 * createTime：2022/1/11 15:09
 * description：开启多个线程批量插入数据
 * 防止多线程情况下数据更新出现异常方式（1--更新字段的时候，添加字段的值，例：balance=balance+100;2加上行的版本号。例：where id =1 and version=1）
 *
 */
public class GenExecutor {
    //开启db线程连接
    static ThreadLocal<Connection> dbConn;

    //构造器连接的初始化
    public GenExecutor() throws ClassNotFoundException {
        Class.forName("com.mysql.jdbc.Driver");
        dbConn = ThreadLocal.withInitial(() -> {
            try {
                return getConnection();
            } catch (SQLException e) {
                throw new RuntimeException(e.getMessage());
            }
        });
    }

    //创建表初始化
//    private void createTable() {
//
//    }

    public static Connection getConnection() throws SQLException {
        return DriverManager.getConnection("xxx",
                "xxx", "xxx");
    }

    //需要执行的sql
    private static void execute(String sql) throws SQLException {
        dbConn.get().createStatement().execute(sql);
    }

    //静态类实现callable
    static class Worker implements Callable {

        static AtomicInteger counter = new AtomicInteger(0);
        private final int id;
        private final int num;
        private final int bucket;

        public Worker(int id, int num, int bucket) {
            this.id = id;
            this.num = num;
            this.bucket = bucket;
        }

        @Override
        public Object call() throws Exception {
            System.out.format("run worker %d\n", id);

            // 10 threads
            int buckets = this.num / this.bucket;

            // 1 user -> 10 post
            int totalUser = this.num / 10;
            int userBuckets = buckets / 10;

            for (int j = 0; j < userBuckets; j++) {
                RowGen rowGen = new RowGen();
                try {
                    // create users
                    String sql = rowGen.genUserBatch(this.bucket);
                    // create posts
                    // 1000 users -> 100,000 posts
                    int userStart = this.id * totalUser + j * 1000;
                    int userEnd = userStart + 1000;
                    //这里应该控制一下批次
                    for (int i = 0; i < 10; i++) {
                        String postSql = rowGen.getBatchPost(1000, userStart, userEnd);
                        execute(postSql);
                    }
                    execute(sql);
                    System.out.format("finished %d/%d\n", counter.incrementAndGet(), buckets);
                } catch (SQLException throwables) {
                    throwables.printStackTrace();
                }
            }
            return null;
        }

    }

    public static Stream<Pair<Integer, Integer>> batch(Integer start, Integer end, Integer batch) {
        return IntStream.iterate(0, o -> o + 1)
                .limit((long) Math.ceil((end - start) / batch))
                //不可变的set value left,right
                .mapToObj(i -> Pair.of(start + i * batch, Math.min(start + (i + 1) * batch, end)));
    }

    public void run(Integer num, Integer bucket) {
        //使用数据库自动建表吧
        //this.createTable();
        //固定的线程数10
        ExecutorService exec = Executors.newFixedThreadPool(10);
        List<Future> futures = (List<Future>) Stream.iterate(0, x -> x + 1)
                .limit(10)
                .map(i -> exec.submit(new Worker(i, num / 10, bucket)))
                .collect(Collectors.toList());//                .forEach(future -> {
//                    System.out.println(future);
//                    try {
//                        future.get();
//                    } catch (InterruptedException | ExecutionException e) {
//                        e.printStackTrace();
//                    }
//                });
        for (Future ftItem : futures) {
            System.out.println(ftItem);
            try {
                ftItem.get();
            } catch (Exception e) {
                throw new RuntimeException(e.getMessage());
            }
        }
    }

    @Test
    public void start() throws ClassNotFoundException {
        GenExecutor gen = new GenExecutor();
        //class worker is div 10 10个og_user

        //写入1000w数据，大概花费20分钟
        gen.run(100000000, 1000);
    }


}
